
from magpie import Magpie
import os
import re
from StringDealer import StringDealer

global_tags = {'字符串', '链表', '快速幂', '回溯', '矩阵快速幂', '逆元', '递归', '字典树', '队列', '哈希', '位运算', '动态规划', '堆', '双指针', '模拟', '数学', '后缀数组', '树', 'bfs', '穷举', '拓扑排序', '二分', '高级结构', '优先队列', '计算几何', '欧拉函数', '三分', '高级算法', '查找', '语法题', '分治', '思维', '图', '排序', '贪心', '前缀和', '搜索', 'dfs', '栈', '复杂度', '数组'}
the_name = "500-20p"

def getLabels():
    base_file_name = "data/"
    files = os.listdir(base_file_name)

    all_tags = set()
    for file in files:
        print(file)
        with open(base_file_name+file,"r",encoding="utf8") as f:

            line = f.readline()
            while line:
                # ss = line.split('\t')
                #
                #
                # print(len(ss))
                # # print(ss[0])
                # # print(ss[1])

                tags = line.split('\t')[1].split(" ")

                for tag in tags:
                    all_tags.add(tag.strip())

                line = f.readline()

    return all_tags


# 训练模型&保存
def trainNSave(dir_name:str):

    #######################
    ### 1. 训练模型并保存 ###
    #######################

    # 先把原来的文件删了
    embeddings_file_name = 'model/embeddings/'+the_name
    scaler_file_name = 'model/scaler/'+the_name
    model_name = 'model/model/'+the_name+'.h5'

    try:
        os.remove(embeddings_file_name)
    except:
        print("删除完成")

    try:
        os.remove(scaler_file_name)
    except:
        print("删除完成")

    try:
        os.remove(model_name)
    except:
        print("删除完成")

    magpie = Magpie()

    # 训练词向量
    magpie.init_word_vectors(dir_name, vec_dim=500)

    # 训练模型
    magpie.train(dir_name, list(global_tags), epochs=15)

    # 保存模型
    magpie.save_word2vec_model(embeddings_file_name)
    magpie.save_scaler(scaler_file_name, overwrite=True)
    magpie.save_model(model_name)

    return magpie


# 加载已有模型,并返回
def load():
    magpie = Magpie(
        keras_model='model/model/' + the_name + '.h5',
        word2vec_model='model/embeddings/' + the_name,
        scaler='model/scaler/' + the_name,
        labels=list(global_tags)
    )

    return magpie

#
def loadNPredict():
    ##################
    ### 2. 加载模型 ###
    #################
    magpie = Magpie(
        keras_model='model/model/' + the_name + '.h5',
        word2vec_model='model/embeddings/' + the_name,
        scaler='model/scaler/' + the_name,
        labels=list(global_tags)
    )

    # 测试
    test_cases = os.listdir("predict/")

    for case in test_cases:
        file_name = "predict/" + case
        print(file_name)

        res = magpie.predict_from_file(file_name)
        print(res)

#
def loadNPredict2():
    ##################
    ### 2. 加载模型 ###
    #################
    magpie = Magpie(
        keras_model='model/model/' + the_name + '.h5',
        word2vec_model='model/embeddings/' + the_name,
        scaler='model/scaler/' + the_name,
        labels=list(global_tags)
    )

    # 测试
    test_cases = os.listdir("predict/")

    for case in test_cases:
        file_name = "predict/" + case
        print(file_name)

        with open(file_name,'r',encoding="utf-8") as f:
            content = f.read();
            content = re.sub("\\s+",",",content)
            content = StringDealer().split(content)
            res = magpie.predict_from_text(content)
            print(res)

#
def loadNPredict3(problem:str):
    ##################
    ### 2. 加载模型 ###
    #################
    magpie = Magpie(
        keras_model='model/model/' + the_name + '.h5',
        word2vec_model='model/embeddings/' + the_name,
        scaler='model/scaler/' + the_name,
        labels=list(global_tags)
    )

    # 测试
    problem = StringDealer().split(problem)
    problem = re.sub("\\s+", ",", problem)
    res = magpie.predict_from_text(problem)
    print(res)
    # print(type(res)) #返回值是list

if __name__ == '__main__':
    print("输入:\n"
          "1:训练模型\n"
          "2:加载模型,并预测predict/下的文本\n"
          "3:加载模型,并预测predict/下的文本(文本未分词)\n"
          "4:加载模型,并输入想要预测的文本\n")
    ins = int(input())

    if ins == 1:
        trainNSave("segment/")
    elif ins == 2:
        loadNPredict()
    elif ins == 3:
        loadNPredict2()
    elif ins == 4:
        print("请输入想要预测的文本")
        p = input()
        loadNPredict3(p)
    else:
        print("bye")





